#!/usr/bin/env python
# -*- coding:utf-8 -*-
import matrixslow as ms
import numpy as np
from matrixslow.core.node import Variable

activate_funs = {
    'ReLU': ms.ops.LeakyReLU,
    'Logistic': ms.ops.Logistic
}


def fc(input, input_size, output_size, activation):
    """
    :param input: 输入向量
    :param input_size: 输入向量的维度
    :param size: 神经元个数，即输出个数（输出向量的维度）
    :param activation: 激活函数类型
    :return: 输出向量
    """
    weights = Variable(dim=(output_size, input_size), init=True, trainable=True)
    bias = Variable(dim=(output_size, 1), init=True, trainable=True)
    affine = ms.ops.Add(ms.ops.MatMul(weights, input), bias)

    return activate_funs.get(activation, lambda x: x)(affine)


def conv(feature_maps, input_shape, kernels, kernel_shape, activation):
    """
    :param feature_maps: 数组，包含多个输入特征图，它们应该是值为同形状的矩阵的节点
    :param input_shape: tuple ，包含输入特征图的形状（宽和高）
    :param kernels: 整数，卷积层的卷积核数量
    :param kernel_shape: tuple ，卷积核的形状（宽和高）
    :param activation: 激活函数类型
    :return: 数组，包含多个输出特征图，它们是值为同形状的矩阵的节点
    """
    # 与输入同形状的全 1 矩阵
    ones = Variable(input_shape, init=False, trainable=False)
    ones.set_value(np.mat(np.ones(input_shape)))

    outputs = []
    for i in range(kernels):

        channels = []
        for fm in feature_maps:
            # 每个channel用的卷积参数不一样，假如图像size是3×224×224，kernel_size=5，
            # 则kernel大小应该是3×5×5，每一层的kernel参数都不一样
            # 这里可以精进一点，直接使用三维张量进行计算，速度肯定会更快，但matrixslow只实现了二维矩阵
            kernel = Variable(kernel_shape, init=True, trainable=True)
            conv = ms.ops.Convolve(fm, kernel)
            channels.append(conv)

        channles = ms.ops.Add(*channels)
        bias = ms.ops.ScalarMultiply(Variable((1, 1), init=True, trainable=True), ones)
        affine = ms.ops.Add(channles, bias)

        outputs.append(activate_funs.get(activation, lambda x: x)(affine))

    assert len(outputs) == kernels
    return outputs


def pooling(feature_maps, kernel_shape, stride):
    """
    :param feature_maps: 数组，包含多个输入特征图，它们应该是值为同形状的矩阵的节点
    :param kernel_shape: tuple ，池化核（窗口）的形状（宽和高）
    :param stride: tuple ，包含横向和纵向步幅
    :return: 数组，包含多个输出特征图，它们是值为同形状的矩阵的节点
    """
    outputs = []
    for fm in feature_maps:
        outputs.append(ms.ops.MaxPooling(fm, size=kernel_shape, stride=stride))

    return outputs
